import torch
import torch.nn.functional as F
import cv2
import numpy as np
import matplotlib.pyplot as plt

# 加载图像并转换为灰度图
image = cv2.imread('touxiang.jpg', cv2.IMREAD_GRAYSCALE)

# 确保图像加载成功
if image is None:
    print("Error loading image")
    exit()

# 转换为 PyTorch 张量，并进行归一化
image_tensor = torch.tensor(image, dtype=torch.float32) / 255.0
image_tensor = image_tensor.unsqueeze(0).unsqueeze(0)  # 添加批次和通道维度

# Sobel 卷积核
sobel_x = torch.tensor([[1, 0, -1],
                        [2, 0, -2],
                        [1, 0, -1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0)  # 形状 [1, 1, 3, 3]
sobel_y = torch.tensor([[1, 2, 1],
                        [0, 0, 0],
                        [-1, -2, -1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0)  # 形状 [1, 1, 3, 3]

# 使用 PyTorch 卷积计算水平和垂直梯度
grad_x = F.conv2d(image_tensor, sobel_x, padding=1)
grad_y = F.conv2d(image_tensor, sobel_y, padding=1)

# 计算梯度幅度（边缘强度）
grad_mag = torch.sqrt(grad_x**2 + grad_y**2)

# 去除多余的维度，并将结果缩放至 [0, 255] 范围
grad_mag = grad_mag.squeeze().numpy()  # 转为 NumPy 数组
grad_mag = np.uint8(np.clip(grad_mag * 255.0, 0, 255))  # 转为 8 位图像

# 显示原图和边缘图
plt.subplot(1, 2, 1)
plt.imshow(image, cmap='gray')
plt.title("Original Image")
plt.axis('off')

plt.subplot(1, 2, 2)
plt.imshow(grad_mag, cmap='gray')
plt.title("Sobel Edge Detection")
plt.axis('off')

plt.show()
